package net.pragyah.scalby.inference.exact


import junit.framework.TestCase
import junit.framework.Assert._
 
import net.pragyah.scalby._

import net.pragyah.scalgorithms.graphs._

class TestMessagePassing_Tree  extends TestCase{
  
	val h = new Variable[String]("H",Set[String]()+"h1"+"h2",None) 
	val b = new Variable[String]("B",Set[String]()+"b1"+"b2",None) 
	val l = new Variable[String]("L",Set[String]()+"l1"+"l2",None) 
	val c = new Variable[String]("C",Set[String]()+"c1"+"c2",None) 
	
	val H = Vertex[Variable[String]](h);
	val B = Vertex[Variable[String]](b);
	val L = Vertex[Variable[String]](l);
	val C = Vertex[Variable[String]](c);
	
	val P = ProbabilityDistribution[String]()
	
	val graph:Graph[Variable[String]] = Graph[Variable[String]](Nil,true)

  
    def testInitialTree{
          

        
        val algo = new MessagePassing_Tree[String](graph,P,H)
        
        algo.initial_tree;
        
        
        val hTag = H.tag.asInstanceOf[VertexTag[String]]
        assertEquals(0.2f,hTag.P("h1"))
        assertEquals(0.8f,hTag.P("h2"))
        assertEquals(1f,hTag.lambda("h1"))
        assertEquals(1f,hTag.lambda("h2"))
        assertEquals(0.2f,hTag.pi("h1"))
        assertEquals(0.8f,hTag.pi("h2"))
        assert(h() == None)

        val bTag = B.tag.asInstanceOf[VertexTag[String]]
        assertEquals(0.09f,bTag.P("b1"))
        assertEquals(910,Math.round(bTag.P("b2")*1000))//this should be 0.91 but comes out to 0.9099997
        assertEquals(1f,bTag.lambda("b1"))
        assertEquals(1f,bTag.lambda("b2"))
        assertEquals(0.09f,bTag.pi("b1"))
        assertEquals(910,Math.round(bTag.pi("b2")*1000).intValue) //this should be 0.91 but comes out to 0.9099997
        assert(b()== None)
        
        val lTag = L.tag.asInstanceOf[VertexTag[String]]
        assertEquals(64,Math.round(lTag.P("l1")*100000))
        assertEquals(0.99936f,lTag.P("l2"))
        assertEquals(1f,lTag.lambda("l1"))
        assertEquals(1f,lTag.lambda("l2"))
        assertEquals(64,Math.round(lTag.pi("l1")*100000))
        assertEquals(0.99936f,lTag.pi("l2"))
        assert(l() == None)
        
        val cTag = C.tag.asInstanceOf[VertexTag[String]]
        assertEquals(2037,Math.round(cTag.P("c1")*100000))
        assertEquals(97963,Math.round(cTag.P("c2")*100000))
        assertEquals(1f,cTag.lambda("c1"))
        assertEquals(1f,cTag.lambda("c2"))
        assertEquals(2037,Math.round(cTag.pi("c1")*100000))
        assertEquals(97963,Math.round(cTag.pi("c2")*100000))
        assert(c() == None)
        
        
        
        val hbTag = graph.getEdge(H,B).get.tag.asInstanceOf[EdgeTag[String]]
        assertEquals(1.0f,hbTag.lambda("h1"))
        assertEquals(1.0f,hbTag.lambda("h2"))
        assertEquals(0.2f,hbTag.pi("h1"))
        assertEquals(0.8f,hbTag.pi("h2"))
        
        
        val hlTag = graph.getEdge(H,L).get.tag.asInstanceOf[EdgeTag[String]]
        assertEquals(1.0f,hlTag.lambda("h1"))
        assertEquals(1.0f,hlTag.lambda("h2"))
        assertEquals(0.2f,hlTag.pi("h1"))
        assertEquals(0.8f,hlTag.pi("h2"))
        
        val lcTag = graph.getEdge(L,C).get.tag.asInstanceOf[EdgeTag[String]]
        assertEquals(1.0f,lcTag.lambda("l1"))
        assertEquals(1.0f,lcTag.lambda("l2"))
        assertEquals(64,Math.round(lcTag.pi("l1")*100000))
        assertEquals(99936,Math.round(lcTag.pi("l2")*100000))

    }
    
    def test_B_b1{
        val algo = new MessagePassing_Tree[String](graph,P,H)

        algo.initial_tree;

        algo.update_tree(b,"b1") //Example 3.4 .. chapter 3 .. Neopolitan

        val hTag = H.tag.asInstanceOf[VertexTag[String]]
        assertEquals(5556,Math.round(hTag.P("h1")*10000).intValue)
        assertEquals(4444,Math.round(hTag.P("h2")*10000).intValue)
        assertEquals(0.25f,hTag.lambda("h1"))
        assertEquals(0.05f,hTag.lambda("h2"))
        assertEquals(0.2f,hTag.pi("h1"))
        assertEquals(0.8f,hTag.pi("h2"))
        assert(h() == None)

        val bTag = B.tag.asInstanceOf[VertexTag[String]]
        assertEquals(1.0f,bTag.P("b1"))
        assertEquals(0f,bTag.P("b2")) //this should be 0.91 but comes out to 0.9099997
        assertEquals(1f,bTag.lambda("b1"))
        assertEquals(0f,bTag.lambda("b2"))
        assertEquals(1.0f,bTag.pi("b1"))
        assertEquals(0f,bTag.pi("b2")) //this should be 0.91 but comes out to 0.9099997
        assertEquals(Some("b1"),b())
        
        val lTag = L.tag.asInstanceOf[VertexTag[String]]
        assertEquals(169,Math.round(lTag.P("l1")*100000))
        assertEquals(99831,Math.round(lTag.P("l2")*100000))
        assertEquals(1f,lTag.lambda("l1"))
        assertEquals(1f,lTag.lambda("l2"))
        assertEquals(15,Math.round(lTag.pi("l1")*100000))
        assertEquals(8985,Math.round(lTag.pi("l2")*100000))
        assert(l() == None)
        
        val cTag = C.tag.asInstanceOf[VertexTag[String]]
        assertEquals(21,Math.round(cTag.P("c1")*1000))
        assertEquals(979,Math.round(cTag.P("c2")*1000))
        assertEquals(1f,cTag.lambda("c1"))
        assertEquals(1f,cTag.lambda("c2"))
        assertEquals(189,Math.round(cTag.pi("c1")*100000))
        assertEquals(8811,Math.round(cTag.pi("c2")*100000))
        assert(c() == None)
        
        
        
        val hbTag = graph.getEdge(H,B).get.tag.asInstanceOf[EdgeTag[String]]
        assertEquals(0.25f,hbTag.lambda("h1"))
        assertEquals(0.05f,hbTag.lambda("h2"))
        assertEquals(0.2f,hbTag.pi("h1"))
        assertEquals(0.8f,hbTag.pi("h2"))
        
        
        val hlTag = graph.getEdge(H,L).get.tag.asInstanceOf[EdgeTag[String]]
        assertEquals(1.0f,hlTag.lambda("h1"))
        assertEquals(1.0f,hlTag.lambda("h2"))
        assertEquals(0.05f,hlTag.pi("h1"))
        assertEquals(4,Math.round(hlTag.pi("h2")*100).intValue)
        
        val lcTag = graph.getEdge(L,C).get.tag.asInstanceOf[EdgeTag[String]]
        assertEquals(1.0f,lcTag.lambda("l1"))
        assertEquals(1.0f,lcTag.lambda("l2"))
        assertEquals(15,Math.round(lcTag.pi("l1")*100000))
        assertEquals(8985,Math.round(lcTag.pi("l2")*100000))
        

        
    }
    
    def test_B_b1_C_c1{

        val algo = new MessagePassing_Tree[String](graph,P,H)

        algo.initial_tree;

        algo.update_tree(b,"b1") //Example 3.4 .. chapter 3 .. Neopolitan
        algo.update_tree(c,"c1") //Example 3.5 .. chapter 3 .. Neopolitan

        val hTag = H.tag.asInstanceOf[VertexTag[String]]
        assertEquals(57569,Math.round(hTag.P("h1")*100000).intValue)
        assertEquals(42431,Math.round(hTag.P("h2")*100000).intValue)
        assertEquals(544,Math.round(hTag.lambda("h1")*100000).intValue)
        assertEquals(100,Math.round(hTag.lambda("h2")*100000).intValue)
        assertEquals(0.2f,hTag.pi("h1"))
        assertEquals(0.8f,hTag.pi("h2"))
        assert(h() == None)

        val bTag = B.tag.asInstanceOf[VertexTag[String]]
        assertEquals(1.0f,bTag.P("b1"))
        assertEquals(0f,bTag.P("b2")) //this should be 0.91 but comes out to 0.9099997
        assertEquals(1f,bTag.lambda("b1"))
        assertEquals(0f,bTag.lambda("b2"))
        assertEquals(1.0f,bTag.pi("b1"))
        assertEquals(0f,bTag.pi("b2")) //this should be 0.91 but comes out to 0.9099997
        assertEquals(Some("b1"),b())
        
        val lTag = L.tag.asInstanceOf[VertexTag[String]]
        assertEquals(4830,Math.round(lTag.P("l1")*100000))
        assertEquals(95170,Math.round(lTag.P("l2")*100000))
        assertEquals(0.6f,lTag.lambda("l1"))
        assertEquals(0.02f,lTag.lambda("l2"))
        assertEquals(15,Math.round(lTag.pi("l1")*100000))
        assertEquals(8985,Math.round(lTag.pi("l2")*100000))
        assert(l() == None)
        
        val cTag = C.tag.asInstanceOf[VertexTag[String]]
        assertEquals(1.0f,cTag.P("c1"))
        assertEquals(0.0f,cTag.P("c2"))
        assertEquals(1.0f,cTag.lambda("c1"))
        assertEquals(0.0f,cTag.lambda("c2"))
        assertEquals(1.0f,cTag.pi("c1"))
        assertEquals(0.0f,cTag.pi("c2"))
        assertEquals(Some("c1"),c())
        
        
        
        val hbTag = graph.getEdge(H,B).get.tag.asInstanceOf[EdgeTag[String]]
        assertEquals(0.25f,hbTag.lambda("h1"))
        assertEquals(0.05f,hbTag.lambda("h2"))
        assertEquals(0.2f,hbTag.pi("h1"))
        assertEquals(0.8f,hbTag.pi("h2"))
        
        
        val hlTag = graph.getEdge(H,L).get.tag.asInstanceOf[EdgeTag[String]]
        assertEquals(2174,Math.round(hlTag.lambda("h1")*100000).intValue)
        assertEquals(2003,Math.round(hlTag.lambda("h2")*100000).intValue)
        assertEquals(0.05f,hlTag.pi("h1"))
        assertEquals(4,Math.round(hlTag.pi("h2")*100).intValue)
        
        val lcTag = graph.getEdge(L,C).get.tag.asInstanceOf[EdgeTag[String]]
        assertEquals(0.6f,lcTag.lambda("l1"))
        assertEquals(0.02f,lcTag.lambda("l2"))
        assertEquals(15,Math.round(lcTag.pi("l1")*100000).intValue)
        assertEquals(8985,Math.round(lcTag.pi("l2")*100000).intValue)
        

    
    }    
    
    /// SETUP 
    
	    val all = H::B::L::C::Nil 
	    graph ++  all//adding all the vertices
     
     
        graph.addEdge(H,B)
        graph.addEdge(H,L)
        graph.addEdge(L,C)
        
        P(h,"h1",Map[Variable[String],String](),0.2f);
        P(h,"h2",Map[Variable[String],String](),0.8f);
        
        P(b,"b1",Map[Variable[String],String]((h,"h1")),0.25f);
        P(b,"b2",Map[Variable[String],String]((h,"h1")),0.75f);
        P(b,"b1",Map[Variable[String],String]((h,"h2")),0.05f);
        P(b,"b2",Map[Variable[String],String]((h,"h2")),0.95f);
        
        P(l,"l1",Map[Variable[String],String]((h,"h1")),0.003f);
        P(l,"l2",Map[Variable[String],String]((h,"h1")),0.997f);
        P(l,"l1",Map[Variable[String],String]((h,"h2")),0.00005f);
        P(l,"l2",Map[Variable[String],String]((h,"h2")),0.99995f);
        
        P(c,"c1",Map[Variable[String],String]((l,"l1")),0.6f);
        P(c,"c2",Map[Variable[String],String]((l,"l1")),0.4f);
        P(c,"c1",Map[Variable[String],String]((l,"l2")),0.02f);
        P(c,"c2",Map[Variable[String],String]((l,"l2")),0.98f);
      

}
